"""
@File    : gdal_backend.py.py
@Time    : 2025/7/15 下午8:27
@Author  : xiashuobad
@Desc    : 

/**
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""
# rsio/core/backends/gdal_backend.py
import numpy as np
from pathlib import Path
from osgeo import gdal, osr
from typing import Literal, Union, Type
from .pil_backend import PILImage, ImageBackend
from ..utils.exceptions import *
from ..utils import (catch_exception, require_attr_not_empty, type_check, PaletteConverter, DataTypeConverter,
                     ArrayShapeConverter, require_python_library, GeoConverter)


class GDALImage(PILImage, backend_name='gdal'):
    open_mode = Literal['r', 'u', 'w']

    def __init__(self, file_path: str, mode: open_mode = 'r', **kwargs):
        super().__init__(file_path, mode, **kwargs)

    @property
    @require_attr_not_empty('img_data')
    def size(self) -> tuple[int, ...]:
        return self.img_data.RasterXSize, self.img_data.RasterYSize

    @property
    @require_attr_not_empty('img_data')
    def shape(self) -> tuple[int, ...]:
        return self.channel, self.height, self.width

    @property
    @require_attr_not_empty('img_data')
    def width(self) -> int:
        return self.img_data.RasterXSize

    @property
    @require_attr_not_empty('img_data')
    def height(self) -> int:
        return self.img_data.RasterYSize

    @property
    @require_attr_not_empty('img_data')
    def channel(self) -> int:
        return self.img_data.RasterCount

    @property
    @require_attr_not_empty('img_data')
    def dtype(self) -> Type[np.number]:
        return DataTypeConverter.gdal_type_to_np(self.img_data.GetRasterBand(1).DataType)

    @property
    @require_attr_not_empty('img_data')
    def geo_bounds(self) -> tuple[float, float, float, float]:
        left, top = self.geo_transform[0], self.geo_transform[3]
        right, bottom = GeoConverter.pixel_to_geo(self.width, self.height, self.geo_transform)
        return left, top, right, bottom

    @property
    @require_attr_not_empty('img_data')
    def nodata(self) -> int | float:
        return self.img_data.GetRasterBand(1).GetNoDataValue()

    @nodata.setter
    @require_attr_not_empty('img_data')
    @type_check
    def nodata(self, nodata: int | float):
        with catch_exception(ImageProcessError('gdal设置nodata失败！')):
            for i in range(1, self.channel + 1):
                self.img_data.GetRasterBand(i).SetNoDataValue(nodata)

    def set_nodata(self, nodata: int | float) -> 'GDALImage':
        self.nodata = nodata
        return self

    @property
    @require_attr_not_empty('img_data')
    def geo_transform(self) -> tuple[float, ...]:
        return self.img_data.GetGeoTransform()

    @geo_transform.setter
    @require_attr_not_empty('img_data')
    @type_check
    def geo_transform(self, geo_transform: tuple[float, ...]):
        self.img_data.SetGeoTransform(GeoConverter.geo_transform2gdal(geo_transform))

    def set_geo_transform(self, geo_transform: tuple[float, ...]) -> 'GDALImage':
        self.geo_transform = geo_transform
        return self

    @property
    @require_attr_not_empty('img_data')
    def geo_projection(self) -> str:
        return self.img_data.GetProjection()

    @geo_projection.setter
    @require_attr_not_empty('img_data')
    @type_check
    def geo_projection(self, geo_projection: str | int | osr.SpatialReference):
        with catch_exception(ImageProcessError('gdal 设置投影失败！')):
            self.img_data.SetProjection(GeoConverter.geo_projection2wkt(geo_projection))

    def set_geo_projection(self, geo_projection: str | int | osr.SpatialReference) -> 'GDALImage':
        self.geo_projection = geo_projection
        return self

    @property
    @require_attr_not_empty('img_data')
    def img_array(self) -> np.ndarray:
        with catch_exception(ImageReadError()):
            return self.img_data.ReadAsArray()

    @img_array.setter
    @type_check
    def img_array(self, img_array: np.ndarray):
        img_array = ArrayShapeConverter.array_shape_to_gdal(img_array)
        if self.img_data is None:
            driver = gdal.GetDriverByName('MEM')
            self._img_data = driver.Create(
                '',
                img_array.shape[-1],
                img_array.shape[-2],
                self.array_channels(img_array),
                DataTypeConverter.np_type_to_gdal(img_array.dtype)
            )
        self._img_data.WriteArray(img_array)

    def set_img_array(self, array: np.ndarray, ) -> 'GDALImage':
        self.img_array = array
        return self

    def _checked_mode(self, mode: open_mode) -> open_mode:
        if mode not in self.open_mode.__args__:
            raise ImageModeError(f'不支持的打开模式：{mode}, 请使用以下模式：{self.open_mode.__args__}！')
        return mode

    @property
    def img_data(self) -> gdal.Dataset:
        return self._img_data

    @img_data.setter
    @type_check
    def img_data(self, img_data: gdal.Dataset):
        driver = gdal.GetDriverByName('MEM')
        self._img_data = driver.CreateCopy('', img_data)

    def set_img_data(self, img_data: gdal.Dataset) -> 'GDALImage':
        self.img_data = img_data
        return self

    @property
    @require_attr_not_empty('img_data')
    def palette(self) -> gdal.ColorTable:
        return self.img_data.GetRasterBand(1).GetColorTable()

    @palette.setter
    @require_attr_not_empty('img_data')
    @type_check
    def palette(self, palette: Union[list, tuple, dict, bytes, np.ndarray, gdal.ColorTable]):
        with catch_exception(ImageProcessError('gdal设置palette失败！')):
            palette: gdal.ColorTable = PaletteConverter.convert('gdal', palette)
            for i in range(1, self.channel + 1):
                self.img_data.GetRasterBand(i).SetColorTable(palette)

    def set_palette(self, palette: Union[list, tuple, dict, bytes, np.ndarray, gdal.ColorTable]) -> 'GDALImage':
        self.palette = palette
        return self

    def _open(self):
        if self.mode != 'w':
            self._img_data = gdal.Open(self.file_path, {'r': gdal.GA_ReadOnly, 'u': gdal.GA_Update}[self.mode])
            if not self._img_data:
                raise ImageOpenError(f"gdal打开文件失败: {self.file_path}")

    @require_attr_not_empty('img_data')
    @ImageBackend.require_write_mode
    def save(self, save_path: str = None):
        with catch_exception(ImageWriteError(f'gdal 保存图像失败！')):
            if self.img_data.GetFileList() is None or (save_path and save_path != self.file_path):
                save_path = save_path or self.file_path
                Path(save_path).resolve().parent.mkdir(parents=True, exist_ok=True)
                driver = gdal.GetDriverByName('GTiff')
                driver.CreateCopy(save_path, self.img_data)
            else:
                self.img_data.FlushCache()

    def close(self):
        self._img_data = None

    @require_attr_not_empty('img_data')
    @type_check
    def resize(self,
               target_size: tuple[int, int] | int,
               keep_ratio: bool = False,
               pad: int = None,
               resample: int = None,
               box: list[int] | tuple[int, int, int, int] = None,
               **kwargs
               ) -> 'GDALImage':
        if isinstance(target_size, int):
            target_size = (target_size, target_size)
        if target_size == self.size:
            return self
        img_data = self.img_data
        if box:
            left, top, right, bottom = box
            img_data = self.crop(left, top, right - left, bottom - top).img_data
        # 原图无坐标系时将坐标轴映射到正方向,防止gdal.Warp报错
        img_data.GetGeoTransform() == (0, 1, 0, 0, 0, 1) and img_data.SetGeoTransform((0, 1, 0, 0, 0, -1))
        # 获取原始尺寸
        original_width, original_height = img_data.RasterXSize, img_data.RasterYSize

        if keep_ratio:
            # 计算保持纵横比的新尺寸
            target_width, target_height = target_size
            # 计算缩放比例
            scale_x = target_width / original_width
            scale_y = target_height / original_height
            scale = min(scale_x, scale_y)
            # 计算新的尺寸
            new_width = round(original_width * scale)
            new_height = round(original_height * scale)
            # 先缩放到保持纵横比的尺寸
            img_data: gdal.Dataset = gdal.Warp(
                '',
                img_data,
                width=new_width,
                height=new_height,
                resampleAlg=resample,
                format='MEM'
            )
            new_image = self.__class__(self.file_path, 'w').set_img_data(img_data)
            if pad is not None:
                # 创建带有padding的目标数据集
                driver = gdal.GetDriverByName('MEM')
                padded_ds: gdal.Dataset = driver.Create(
                    '',
                    target_width,
                    target_height,
                    img_data.RasterCount,
                    img_data.GetRasterBand(1).DataType
                )

                # 计算在目标图像中的位置（居中）
                x_offset = (target_width - new_width) // 2
                y_offset = (target_height - new_height) // 2
                # 将缩放后的图像复制到padding图像的中心
                for i in range(1, img_data.RasterCount + 1):
                    # 先填充padding值
                    padded_band = padded_ds.GetRasterBand(i)
                    padded_band.Fill(pad)
                    # 复制缩放后的图像数据
                    data = img_data.GetRasterBand(i).ReadAsArray()
                    padded_band.WriteArray(data, x_offset, y_offset)

                # 设置地理变换、投影等元数据
                geo_transform = list(img_data.GetGeoTransform())
                # 由于内容向右下偏移，左上角坐标需要相应调整
                geo_transform[0] -= x_offset * geo_transform[1]  # X坐标减小
                geo_transform[3] -= y_offset * geo_transform[5]  # Y坐标增大
                padded_ds.SetGeoTransform(tuple(geo_transform))
                (geo_projection := img_data.GetProjection()) is not None and padded_ds.SetProjection(geo_projection)
                nodata = img_data.GetRasterBand(1).GetNoDataValue()
                palette = PaletteConverter.clone_color_table(img_data.GetRasterBand(1).GetColorTable())
                # 更新图像数据
                new_image._img_data = padded_ds
                # # 设置无数据值
                nodata is not None and new_image.set_nodata(nodata)
                palette is not None and new_image.set_palette(palette)

        else:
            # 直接调整到目标尺寸
            img_data = gdal.Warp(
                '',
                img_data,
                width=target_size[0],
                height=target_size[1],
                resampleAlg=resample,
                format='MEM',
            )
            new_image = self.__class__(self.file_path, 'w').set_img_data(img_data)
        return new_image

    def reverse_rgb(self) -> 'GDALImage':
        array_reversed = self.reversed_rgb(self.img_array, 'gdal')
        self.img_data.WriteArray(array_reversed)
        return self

    @require_attr_not_empty('img_data')
    @type_check
    def read_window_array(self, x_off: int = 0, y_off: int = 0, window_x_size: int = None,
                          window_y_size: int = None,
                          band_indexes: int | list[int] = None, **kwargs):
        if isinstance(band_indexes, int):
            band_indexes = [band_indexes]
        if isinstance(band_indexes, list):
            for i, band_index in enumerate(band_indexes):
                if band_index < 0 or band_index >= self.channel:
                    raise ValueError(f'波段索引值：{band_indexes} 超出了范围：{tuple(range(self.channel))}')
                band_indexes[i] = band_index + 1
        return self.img_data.ReadAsArray(x_off, y_off, window_x_size, window_y_size, band_list=band_indexes,
                                         **kwargs)

    @require_attr_not_empty('img_data')
    @type_check
    def crop(self,
             x_off: int = 0,
             y_off: int = 0,
             window_x_size: int = None,
             window_y_size: int = None
             ) -> 'GDALImage':
        x_off, y_off, window_x_size, window_y_size = self._checked_crop_window(x_off, y_off, window_x_size,
                                                                               window_y_size)
        cropped_ds: gdal.Dataset = gdal.Translate(
            '',
            self.img_data,
            srcWin=[x_off, y_off, window_x_size, window_y_size],
            format='MEM',
        )
        if not cropped_ds:
            raise ImageProcessError('裁剪图片失败！')
        new_image = self.__class__('', 'w').set_img_data(cropped_ds)
        return new_image

    @require_attr_not_empty('img_data')
    @require_python_library('PIL', 'pillow')
    def show(self, title: str = None):
        from PIL import Image
        image_show = Image.fromarray(ArrayShapeConverter.convert('pil', self.img_array))
        if self.palette:
            image_show.putpalette(PaletteConverter.convert('pil', self.palette))
        image_show.show(title)

    def reproject(self, dst_srs_obj: 'int|str|GDALImage', dst_width=None, dst_height=None, num_threads: int = 2):
        if not self.geo_projection:
            raise ImageProcessError(f'图像:{self.file_name} 未定义投影信息，请先设置投影！')
        if isinstance(dst_srs_obj, GDALImage):
            dst_srs_wkt = dst_srs_obj.geo_projection
            dst_width, dst_height = dst_srs_obj.size
        else:
            dst_srs_wkt = GeoConverter.geo_projection2wkt(dst_srs_obj)

        warp_options = gdal.WarpOptions(
            format='MEM',
            srcSRS=self.geo_projection,
            dstSRS=dst_srs_wkt,
            resampleAlg=gdal.GRA_Bilinear,
            multithread=num_threads is not None,  # 是否启用多线程
            warpOptions=dict(NUM_THREADS=num_threads) if num_threads else None,
            width=dst_width,
            height=dst_height,
        )
        # 执行重投影
        self._img_data = gdal.Warp('', self.img_data, options=warp_options)
        return self

    @require_attr_not_empty('img_array')
    def convert_to_rgb(self):
        new_image = GDALImage(
            self.file_path.replace(self.img_suffix, f'_to_rgb{self.img_suffix}'),
            'w',
        )
        array_to_rgb = ArrayShapeConverter.array_shape_to_rgb(self.img_array)
        new_image.set_img_array(array_to_rgb).set_geo_transform(
            self.geo_transform).set_geo_projection(self.geo_projection)

        return new_image

    @staticmethod
    @type_check
    def write_array(
            fp: str,
            array: np.ndarray,
            geo_transform: tuple[float, ...] = None,
            geo_projection: str = None,
            nodata: Union[int, float] = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,
            **kwargs
    ):
        with catch_exception(ImageWriteError(f'gdal 保存图片失败！')):
            array = ArrayShapeConverter.array_shape_to_gdal(array)
            driver = gdal.GetDriverByName('GTiff')
            dtype = DataTypeConverter.np_type_to_gdal(array.dtype)
            dataset: gdal.Dataset = driver.Create(
                fp,
                array.shape[-1],
                array.shape[-2],
                array.shape[0] if len(array.shape) == 3 else 1,
                dtype
            )
            dataset.WriteArray(array, **kwargs)
            geo_transform and dataset.SetGeoTransform(geo_transform)
            geo_projection and dataset.SetProjection(geo_projection)
            nodata and [dataset.GetRasterBand(i).SetNoDataValue(nodata) for i in range(1, array.shape[0] + 1)]
            if palette:
                palette = PaletteConverter.convert('gdal', palette)
                for i in range(1, array.shape[0] + 1):
                    dataset.GetRasterBand(i).SetColorTable(palette)
            dataset.FlushCache()
